import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, Circle, ConnectionPatch

# 图3.1 AH-DRLP算法整体架构图
fig, ax = plt.subplots(1, 1, figsize=(12, 8))
fig.suptitle('图3.1 AH-DRLP算法整体架构', fontproperties='SimHei', fontsize=16)

# 定义模块位置和大小
modules = {
    'env': {'xy': (0.05, 0.4), 'size': (0.15, 0.15), 'label': '环境\n状态 $s_t$', 'color': '#E1F5FE'},
    'state_encoder': {'xy': (0.25, 0.35), 'size': (0.2, 0.25), 'label': '状态编码器\n(CNN/LSTM+Attention)', 'color': '#C8E6C9'},
    'high_level': {'xy': (0.5, 0.6), 'size': (0.2, 0.15), 'label': '高层策略 (Meta-Controller)\nActor-Critic', 'color': '#FFF9C4'},
    'low_level': {'xy': (0.5, 0.3), 'size': (0.2, 0.15), 'label': '底层策略 (Controller)\nActor-Critic', 'color': '#FFCCBC'},
    'action': {'xy': (0.75, 0.3), 'size': (0.1, 0.1), 'label': '动作 $a_t$', 'color': '#F8BBD0'},
    'goal': {'xy': (0.75, 0.7), 'size': (0.15, 0.1), 'label': '最终目标 $g$', 'color': '#D1C4E9'},
}

# 绘制模块
for key, module in modules.items():
    x, y = module['xy']
    w, h = module['size']
    p = FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.05", facecolor=module['color'], edgecolor='black')
    ax.add_patch(p)
    ax.text(x + w/2, y + h/2, module['label'], ha='center', va='center', fontproperties='SimHei', fontsize=9, linespacing=1.5)

# 绘制数据流箭头
flows = [
    # (start_pos, end_pos, label, color, style)
    ((0.2, 0.475), (0.25, 0.475), '', 'black', '-'), # Env -> Encoder
    ((0.45, 0.475), (0.5, 0.475), '编码后状态', 'blue', '-'), # Encoder -> High
    ((0.7, 0.675), (0.75, 0.675), '', 'purple', '-'), # Goal -> High
    ((0.6, 0.6), (0.65, 0.525), '子目标 $g_{sub}$', 'orange', '->'), # High -> Low (via point)
    ((0.45, 0.375), (0.5, 0.375), '编码后状态', 'blue', '-'), # Encoder -> Low
    ((0.7, 0.3), (0.85, 0.3), '', 'red', '-'), # Action -> Env (right)
    ((0.85, 0.3), (0.2, 0.3), '', 'red', '-'), # Env (left) -> Encoder (loop back)
]

# 绘制连接线和箭头
for flow in flows:
    start, end, label, color, style = flow
    if style == '->':
        ax.arrow(start[0], start[1], end[0]-start[0], end[1]-start[1], head_width=0.015, head_length=0.02, fc=color, ec=color, length_includes_head=True)
    else:
        ax.plot([start[0], end[0]], [start[1], end[1]], color=color, linestyle=style)
    if label:
        ax.text((start[0]+end[0])/2, (start[1]+end[1])/2 + 0.02, label, ha='center', fontproperties='SimHei', fontsize=8, color=color)

# 绘制从高层到底层的虚线箭头
ax.annotate('', xy=(0.65, 0.525), xytext=(0.6, 0.6),
            arrowprops=dict(arrowstyle='->', color='orange', linestyle='--', linewidth=1))

ax.text(0.625, 0.56, '子目标 $g_{sub}$', ha='center', va='center', fontproperties='SimHei', fontsize=8, color='orange')

ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.axis('off')
plt.tight_layout()
plt.show()